#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 应用案例.py
# @Author: HCY
# @Date  : 2025/5/17
# @Desc  : 讯飞角色模拟服务应用案例

import os
import time
import json
import gradio as gr
from typing import List, Dict

# SDK引入模型
from dwspark.config import Config
from dwspark.models import CharacterSimulator
# 日志
from loguru import logger

# 加载系统环境变量：SPARKAI_UID、SPARKAI_APP_ID、SPARKAI_API_KEY、SPARKAI_API_SECRET
config = Config()

# 创建角色模拟器实例
simulator = CharacterSimulator(config)

# 存储当前会话的玩家ID和人格ID
current_player_id = None
current_agent_id = None

def create_player(player_name: str, player_desc: str):
    """
    创建玩家
    :param player_name: 玩家名称
    :param player_desc: 玩家描述
    :return: 玩家ID和状态信息
    """
    global current_player_id
    try:
        current_player_id = simulator.create_player(player_name, player_desc)
        return f"创建玩家成功！玩家ID: {current_player_id}", f"已创建玩家: {player_name}", "未创建人格"
    except Exception as e:
        logger.error(f"创建玩家失败: {e}")
        return f"创建玩家失败: {str(e)}", "未创建玩家", "未创建人格"

def create_agent(agent_name: str, agent_identity: str, agent_hobby: str, agent_personality: str):
    """
    创建人格
    :param agent_name: 人格名称
    :param agent_identity: 人格身份
    :param agent_hobby: 人格爱好
    :param agent_personality: 人格性格描述
    :return: 人格ID和状态信息
    """
    global current_player_id, current_agent_id
    if not current_player_id:
        return "请先创建玩家！", "未创建玩家", "未创建人格"
    
    try:
        current_agent_id = simulator.create_agent(
            player_id=current_player_id,
            agent_name=agent_name,
            agent_identity=agent_identity,
            agent_hobby=agent_hobby,
            agent_personality_desc=agent_personality
        )
        return f"创建人格成功！人格ID: {current_agent_id}", f"已创建玩家", f"已创建人格: {agent_name}"
    except Exception as e:
        logger.error(f"创建人格失败: {e}")
        return f"创建人格失败: {str(e)}", f"已创建玩家", "未创建人格"

def chat(chat_query: str, chat_history: List):
    """
    与角色对话
    :param chat_query: 当前用户问题
    :param chat_history: 历史对话
    :return: 更新后的对话历史
    """
    global current_player_id, current_agent_id
    if not current_player_id or not current_agent_id:
        chat_history.append((chat_query, "请先创建玩家和人格！"))
        return "", chat_history
    
    try:
        # 添加用户问题到对话历史
        chat_history.append((chat_query, ""))
        
        # 定义回调函数
        # 用于累积消息文本
        response_text = []
        response_complete = False
        
        def on_message(data):
            nonlocal chat_history, response_text, response_complete
            logger.debug(f"收到消息: {data}")
            
            # 尝试解析字符串消息
            if isinstance(data, str):
                try:
                    data = json.loads(data)
                except json.JSONDecodeError:
                    logger.error(f"无法解析JSON消息: {data}")
                    # 如果无法解析为JSON，直接将其作为文本添加
                    response_text.append(data)
                    chat_history[-1] = (chat_query, ''.join(response_text))
                    return

            # 检查不同的消息格式
            # 格式1: payload.choices.text[]
            if isinstance(data, dict):
                # 检查是否有错误码
                if "code" in data and data["code"] != 0:
                    error_msg = data.get("message", "未知错误")
                    chat_history[-1] = (chat_query, f"对话失败: {error_msg}")
                    response_complete = True
                    return
                
                # 处理payload格式
                if 'payload' in data:
                    payload = data['payload']
                    if isinstance(payload, dict):
                        # 处理choices格式
                        if 'choices' in payload:
                            choices = payload['choices']
                            if isinstance(choices, dict) and 'text' in choices:
                                texts = choices['text']
                                if isinstance(texts, list):
                                    # 获取所有文本内容并拼接
                                    content = ''.join([text.get('content', '') for text in texts if isinstance(text, dict)])
                                    if content:
                                        response_text.append(content)
                                        chat_history[-1] = (chat_query, ''.join(response_text))
                        # 直接处理content字段
                        elif 'content' in payload:
                            content = payload['content']
                            if content:
                                response_text.append(content)
                                chat_history[-1] = (chat_query, ''.join(response_text))
                
                # 直接处理content字段
                elif 'content' in data:
                    content = data['content']
                    if content:
                        response_text.append(content)
                        chat_history[-1] = (chat_query, ''.join(response_text))
                
                # 检查是否完成
                if "status" in data and data["status"] == 2:
                    response_complete = True
            
        def on_error(error_msg):
            nonlocal chat_history, response_complete
            logger.error(f"对话失败: {error_msg}")
            chat_history[-1] = (chat_query, f"对话失败: {error_msg}")
            response_complete = True
            
        def on_close():
            nonlocal response_complete
            logger.info("对话连接已关闭")
            response_complete = True
        
        # 创建新会话
        chat_id = simulator.create_chat(player_id=current_player_id, agent_id=current_agent_id)
        if not chat_id:
            raise Exception("创建会话失败")
            
        # 调用角色模拟器进行对话，添加回调函数
        simulator.chat(
            chat_id=chat_id,
            player_id=current_player_id,
            agent_id=current_agent_id,
            content=chat_query,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close
        )
        
        # 等待响应完成
        max_wait_time = 30  # 最大等待时间，单位秒
        wait_start = time.time()
        
        # 先等待一小段时间，让回调有机会执行
        time.sleep(2)
        
        # 如果响应为空，可能是回调没有被调用
        if not response_text:
            chat_history[-1] = (chat_query, "未收到回复，请检查网络连接或重试")
        
        # 返回清空输入框和更新后的对话历史
        return "", chat_history
    except Exception as e:
        logger.error(f"对话失败: {e}")
        chat_history[-1] = (chat_query, f"对话失败: {str(e)}")
        return "", chat_history

# 构建Gradio界面
with gr.Blocks(title="角色模拟器") as demo:
    gr.Markdown("# 讯飞星火角色模拟器")
    
    # 创建状态显示组件
    status_display1 = gr.Textbox(label="玩家状态", value="未创建玩家", interactive=False)
    status_display2 = gr.Textbox(label="人格状态", value="未创建人格", interactive=False)
    
    with gr.Tab("创建角色"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("## 创建玩家")
                player_name = gr.Textbox(label="玩家名称", placeholder="请输入玩家名称", value="测试玩家")
                player_desc = gr.Textbox(label="玩家描述", placeholder="请输入玩家描述", value="这是一个测试玩家")
                create_player_btn = gr.Button("创建玩家")
                player_result = gr.Textbox(label="创建结果", interactive=False)
                
                create_player_btn.click(
                    fn=create_player,
                    inputs=[player_name, player_desc],
                    outputs=[player_result, status_display1, status_display2]
                )
            
            with gr.Column():
                gr.Markdown("## 创建人格")
                agent_name = gr.Textbox(label="人格名称", placeholder="请输入人格名称", value="小助手")
                agent_identity = gr.Textbox(label="人格身份", placeholder="请输入人格身份", value="AI助手")
                agent_hobby = gr.Textbox(label="人格爱好", placeholder="请输入人格爱好", value="帮助用户解决问题")
                agent_personality = gr.Textbox(label="人格性格描述", placeholder="请输入人格性格描述", value="友善、耐心、专业")
                create_agent_btn = gr.Button("创建人格")
                agent_result = gr.Textbox(label="创建结果", interactive=False)
                
                create_agent_btn.click(
                    fn=create_agent,
                    inputs=[agent_name, agent_identity, agent_hobby, agent_personality],
                    outputs=[agent_result, status_display1, status_display2]
                )
    
    with gr.Tab("角色对话"):
        gr.Markdown("## 与角色对话")
        with gr.Row():
            with gr.Column(scale=2):
                chatbot = gr.Chatbot([], elem_id="chat-box", label="聊天历史")
                chat_input = gr.Textbox(label="输入问题", placeholder="请输入您想问的问题")
                chat_btn = gr.Button("发送")
        
        # 问题样例
        gr.Examples(["你好，请介绍一下你自己", "你有什么爱好？", "你能帮我做什么？"], chat_input)
        
        # 绑定发送按钮点击事件
        chat_btn.click(
            fn=chat,
            inputs=[chat_input, chatbot],
            outputs=[chat_input, chatbot]
        )
        
        # 绑定回车键发送
        chat_input.submit(
            fn=chat,
            inputs=[chat_input, chatbot],
            outputs=[chat_input, chatbot]
        )

if __name__ == "__main__":
    demo.queue().launch()